import backtrader as bt
import talib
import datetime


# 定义策略类
class TechnicalIndicatorsStrategy(bt.Strategy):
    params = (
        ("sma_period", 20),
        ("rsi_period", 14),
    )

    def __init__(self):
        self.close = self.datas[0].close
        self.sma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.sma_period
        )
        self.rsi = bt.talib.RSI(self.datas[0], timeperiod=self.params.rsi_period)
        self.macd, self.macd_signal, _ = bt.talib.MACD(self.datas[0])

    def next(self):
        if not self.position:
            if self.rsi[0] < 30 and self.close[0] > self.sma[0]:
                self.buy()
        else:
            if self.rsi[0] > 70:
                self.sell()


# 创建 Cerebro 实例
cerebro = bt.Cerebro()

# 添加数据
data = bt.feeds.YahooFinanceData(
    dataname="AAPL",
    fromdate=datetime.datetime(2020, 1, 1),
    todate=datetime.datetime(2023, 1, 1),
)
cerebro.adddata(data)

# 添加策略
cerebro.addstrategy(TechnicalIndicatorsStrategy)

# 设置初始资金
cerebro.broker.setcash(100000.0)

# 设置交易手续费
cerebro.broker.setcommission(commission=0.001)

# 运行回测
print("Starting Portfolio Value: %.2f" % cerebro.broker.getvalue())
cerebro.run()
print("Final Portfolio Value: %.2f" % cerebro.broker.getvalue())
